import random
import math


class AntColonyOptimization:
    def __init__(self, num_ants, alpha, beta, rho, q, cities):
        self.num_ants = num_ants
        self.alpha = alpha
        self.beta = beta
        self.rho = rho
        self.q = q
        self.cities = cities
        self.num_cities = len(cities)
        self.pheromone = [[1 for _ in range(self.num_cities)] for _ in range(self.num_cities)]

    def distance(self, city1, city2):
        return math.sqrt((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2)

    def calculate_probability(self, ant_position, visited):
        probabilities = []
        current_city = ant_position
        unvisited_cities = [i for i in range(self.num_cities) if i not in visited]

        for next_city in unvisited_cities:
            pheromone = self.pheromone[current_city][next_city]
            distance = self.distance(self.cities[current_city], self.cities[next_city])
            heuristic = 1 / distance
            probability = (pheromone ** self.alpha) * (heuristic ** self.beta)
            probabilities.append((next_city, probability))

        total_probability = sum(prob for _, prob in probabilities)
        probabilities = [(city, prob / total_probability) for city, prob in probabilities]

        return probabilities

    def select_next_city(self, ant_position, visited):
        probabilities = self.calculate_probability(ant_position, visited)
        random_number = random.random()
        cumulative_probability = 0

        for city, probability in probabilities:
            cumulative_probability += probability
            if random_number < cumulative_probability:
                return city

        return probabilities[-1][0]

    def update_pheromone(self, ant_paths):
        for i in range(self.num_cities):
            for j in range(self.num_cities):
                self.pheromone[i][j] *= (1 - self.rho)

        for ant_path in ant_paths:
            path_length = 0
            for k in range(len(ant_path) - 1):
                from_city = ant_path[k]
                to_city = ant_path[k + 1]
                path_length += self.distance(self.cities[from_city], self.cities[to_city])

            for k in range(len(ant_path) - 1):
                from_city = ant_path[k]
                to_city = ant_path[k + 1]
                self.pheromone[from_city][to_city] += self.q / path_length

    def run(self):
        best_path = None
        best_path_length = float('inf')

        for iteration in range(100):
            ant_paths = []

            for ant in range(self.num_ants):
                visited = [random.randint(0, self.num_cities - 1)]
                current_city = visited[-1]

                while len(visited) < self.num_cities:
                    next_city = self.select_next_city(current_city, visited)
                    visited.append(next_city)
                    current_city = next_city

                visited.append(visited[0])
                ant_paths.append(visited)

                path_length = 0
                for i in range(len(visited) - 1):
                    path_length += self.distance(self.cities[visited[i]], self.cities[visited[i + 1]])

                if path_length < best_path_length:
                    best_path_length = path_length
                    best_path = visited

            self.update_pheromone(ant_paths)

        return best_path, best_path_length


cities = [(0, 0), (1, 2), (3, 1), (5, 3), (2, 5)]
aco = AntColonyOptimization(10, 1, 5, 0.5, 10, cities)
best_path, best_path_length = aco.run()
print("Best path:", best_path)
print("Best path length:", best_path_length)
